package dp.naiveBayes;

import java.io.File;
import java.io.FilenameFilter;
import java.io.IOException;
import java.nio.charset.Charset;
import java.util.List;
import java.util.Map;

import org.apache.mahout.classifier.ClassifierResult;
import org.apache.mahout.classifier.ResultAnalyzer;
import org.apache.mahout.classifier.bayes.Algorithm;
import org.apache.mahout.classifier.bayes.BayesAlgorithm;
import org.apache.mahout.classifier.bayes.BayesParameters;
import org.apache.mahout.classifier.bayes.ClassifierContext;
import org.apache.mahout.classifier.bayes.Datastore;
import org.apache.mahout.classifier.bayes.InMemoryBayesDatastore;
import org.apache.mahout.classifier.bayes.InvalidDatastoreException;
import org.apache.mahout.classifier.bayes.mapreduce.bayes.BayesClassifierDriver;
import org.apache.mahout.common.TimingStatistics;
import org.apache.mahout.common.iterator.FileLineIterable;
import org.apache.mahout.common.nlp.NGrams;

import dp.utils.Utils;

/**
 * 
 * Inspired by TestClassifier.java usage: NBTester [path-to-model(HDFS)]
 * [path-to-test-input(HDFS)] [seq|mapred]
 */
public class NBTester {

	public static void main(String[] args) {

		if (args.length < 2 || args.length > 3) {
			System.err
					.println("zly pocet parametrov: pouzitie: NBTester [path-to-model(HDFS)] [path-to-test-input(HDFS)] [seq|mapred]");
			System.exit(1);
		}

		// TODO: configure by property file
		BayesParameters params = new BayesParameters();
		params.setGramSize(1);
		params.set("verbose", Boolean.toString(true));
		params.setBasePath(args[0]);
		params.set("classifierType", "bayes");
		params.set("dataSource", "hdfs");
		params.set("defaultCat", "unknown");
		params.set("encoding", "UTF-8");
		params.set("alpha_i", "1.0");
		params.set("testDirPath", args[1]);
		// params.set("confusionMatrix", null);

		NBTester tester = new NBTester();
		if (args.length == 3 && "seq".equals(args[2])) {
			tester.classifyBayesSequential(params);
		} else {
			tester.classifyBayesParalel(params);
		}

	}

	public void classifyBayesSequential(BayesParameters params) {
		File dir = new File(params.get("testDirPath"));
		File[] files = dir.listFiles();
		Algorithm algorithm = new BayesAlgorithm();
		Datastore datastore = new InMemoryBayesDatastore(params);
		ClassifierContext classifier = new ClassifierContext(algorithm, datastore);
		try {
			classifier.initialize();
			ResultAnalyzer resultAnalyzer = new ResultAnalyzer(classifier.getLabels(), params.get("defaultCat"));
			TimingStatistics totalStatistics = new TimingStatistics();

			if (files != null) {
				for (File file : files) {
					TimingStatistics operationStats = new TimingStatistics();
					for (String line : new FileLineIterable(new File(file.getPath()), Charset.forName(params
							.get("encoding")), false)) {

						Map<String, List<String>> document = new NGrams(line, Integer.parseInt(params.get("gramSize")))
								.generateNGrams();
						for (Map.Entry<String, List<String>> stringListEntry : document.entrySet()) {
							String correctLabel = stringListEntry.getKey();
							List<String> strings = stringListEntry.getValue();
							TimingStatistics.Call call = operationStats.newCall();
							TimingStatistics.Call outercall = totalStatistics.newCall();
							ClassifierResult classifiedLabel = classifier.classifyDocument(
									strings.toArray(new String[strings.size()]), params.get("defaultCat"));
							call.end();
							outercall.end();
							boolean correct = resultAnalyzer.addInstance(correctLabel, classifiedLabel);

						}
					}
				}
			}
			Utils.log(resultAnalyzer.toString());
		} catch (InvalidDatastoreException e) {
			e.printStackTrace();
		} catch (NumberFormatException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		}

	}

	public void classifyBayesParalel(BayesParameters params) {
		try {
			classifyBayes(params);
		} catch (IOException e) {
			e.printStackTrace();
		}
	}

	public static void classifyBayes(BayesParameters params) throws IOException {
		BayesClassifierDriver.runJob(params);
	}
}
